import psycopg2


def get_schemas_with_size(connection):
    cursor = connection.cursor()
    # 修改后的查询，计算每个模式下所有表的总大小
    query = """
        SELECT n.nspname, SUM(pg_total_relation_size(c.oid)) as size
        FROM pg_class c
        JOIN pg_namespace n ON c.relnamespace = n.oid
        WHERE n.nspname !~ '^pg_' AND n.nspname != 'information_schema'
        GROUP BY n.nspname
        ORDER BY size DESC;
    """
    cursor.execute(query)
    schemas = cursor.fetchall()
    cursor.close()
    return schemas


def print_schemas(schemas):
    print("以下是按大小排序的模式列表：")
    for i, (schema, size) in enumerate(schemas, start=1):
        print(f"{i}. {schema} - {size // (1024 * 1024)} MB")


def delete_schema(connection, schema_index, schemas):
    if 1 <= schema_index <= len(schemas):
        schema_to_delete = schemas[schema_index - 1][0]
        cursor = connection.cursor()
        try:
            # 先删除模式下的所有对象
            cursor.execute(f"DROP SCHEMA {schema_to_delete} CASCADE;")
            connection.commit()
            print(f"成功删除模式 {schema_to_delete}")
        except (Exception, psycopg2.Error) as error:
            print(f"删除模式 {schema_to_delete} 时出错：{error}")
            connection.rollback()
        finally:
            cursor.close()
    else:
        print("输入的编号无效，请输入有效的编号。")


def main():
    try:
        # 建立数据库连接，需要根据实际情况修改
        connection = psycopg2.connect(
            database="net_platfrom",
            user="postgres",
            password="wxxx12345",
            host="192.168.67.56",
            port="5432"
        )

        # 获取模式及其大小
        schemas = get_schemas_with_size(connection)

        # 打印模式列表
        print_schemas(schemas)

        # 提示用户输入要删除的模式编号
        try:
            index = int(input("请输入要删除的模式编号（输入 0 退出）："))
            if index == 0:
                print("操作已取消。")
            else:
                delete_schema(connection, index, schemas)
        except ValueError:
            print("输入无效，请输入一个有效的整数编号。")

    except (Exception, psycopg2.Error) as error:
        print("连接或查询数据库时出错：", error)
    finally:
        if 'connection' in locals() and connection:
            connection.close()


if __name__ == "__main__":
    main()
